# Copyright 2022 The Ip2Region Authors. All rights reserved.
# Use of this source code is governed by a Apache2.0-style
# license that can be found in the LICENSE file.
#
#  Created by luckydog on 2022/6/29.
#  Copyright © 2022年 luckydog. All rights reserved.
#
# 文件来源
# https://github.com/lionsoul2014/ip2region/blob/master/binding/python/xdbSearcher.py

import socket
import struct
from typing import Optional

from app.core.config import TEMP_DIR
from app.core.log import logger

# xdb默认参数
_HEADER_INFO_LENGTH = 256
_VECTOR_INDEX_ROWS = 256
_VECTOR_INDEX_COLS = 256
_VECTOR_INDEX_SIZE = 8
_SEGMENT_INDEX_SIZE = 14

_IPV6_LOCAL_IP = "0:0:0:0:0:0:0:1"
_IPV6_LOCAL_IP_REGION = "0|0|0|内网IP|内网IP"


def ip_to_uint32(ip: str) -> int:
    """
    将IPv4地址字符串转换为32位无符号整数

    Example:
        >>> ip_to_uint32("192.168.1.1")
        3232235777

    Note:
        - 使用大端字节序("!L")进行解包
        - 例如: "1.2.3.4" 转换为 0x01020304 (16909060)

    :param ip: IPv4地址字符串，例如 "192.168.1.1"
    :return: 对应的32位无符号整数表示的IP地址
    """
    _ip = socket.inet_aton(ip)
    return struct.unpack("!L", _ip)[0]


def get_uint32(b: bytes, offset: int) -> int:
    """
    从字节串中指定偏移位置读取4字节数据并转换为无符号32位整数

    Example:
        >>> data = b'\x01\x00\x00\x00\x02\x00\x00\x00'
        >>> get_uint32(data, 0)
        1
        >>> get_uint32(data, 4)
        2

    Note:
        - 使用小端序('I')解包4字节无符号整数
        - 主要用于读取IP地址库中的指针和长度信息
        - 当字节数据不足4字节时返回默认值0

    :param b: 源字节串数据
    :param offset: 起始偏移位置
    :return: 4字节数据转换后的无符号32位整数，如果数据不足4字节则返回0
    """
    if len(b[offset:offset + 4]) == 4:
        return struct.unpack('I', b[offset:offset + 4])[0]
    return 0


def get_uint16(b: bytes, offset: int) -> int:
    """
    从字节串中指定偏移位置读取2字节数据并转换为无符号16位整数

    Example:
        >>> data = b'\x01\x02'
        >>> get_uint16(data, 0)
        513  # 0x0201 = 2*256 + 1 = 513

    Note:
        - 使用小端序方式解析2字节数据
        - 第一个字节为低位字节，第二个字节为高位字节
        - 主要用于读取IP地址库中的长度信息

    解析过程:
        1. b[offset] & 0x000000FF: 获取第一个字节的值(低位字节)
        2. b[offset + 1] << 8: 获取第二个字节的值并向左移8位(高位字节)
        3. 两者进行按位或运算得到最终结果

    :param b: 源字节串数据
    :param offset: 起始偏移位置
    :return: 2字节数据转换后的无符号16位整数
    """
    return (b[offset] & 0x000000FF) | (b[offset + 1] << 8)


def is_ipv4(ip: str) -> bool:
    """
    验证字符串是否为有效的IPv4地址格式

    Example:
        >>> is_ipv4("192.168.1.1")
        True
        >>> is_ipv4("256.1.1.1")
        False
        >>> is_ipv4("192.168.1")
        False

    Note:
        - 验证规则：
          1. 必须包含且仅包含3个点号，分为4个部分
          2. 每部分必须是数字
          3. 每部分长度不能超过3位
          4. 每部分数值必须在0-255范围内

    :param ip: 待验证的IP地址字符串
    :return: 如果是有效的IPv4地址返回True，否则返回False
    """
    ip_parts = ip.split(".")

    # 检查是否恰好有4个部分
    if len(ip_parts) != 4:
        return False

    # 检查每个部分是否有效
    for part in ip_parts:
        # 检查是否为数字
        if not part.isdigit():
            return False
        # 检查长度（0-255最多3位数字）
        if len(part) > 3:
            return False
        # 检查数值范围
        if int(part) > 255:
            return False
    return True


def is_ipv6(ip: str) -> bool:
    """
    验证字符串是否为有效的IPv6地址格式

    Example:
        >>> is_ipv6("2001:0db8:85a3:0000:0000:8a2e:0370:7334")
        True
        >>> is_ipv6("2001:db8:85a3::8a2e:370:7334")
        True
        >>> is_ipv6("::1")
        True
        >>> is_ipv6("::ffff:192.168.1.1")
        True

    Note:
        - IPv6地址由8组十六进制数组成，每组1-4个字符，用冒号分隔
        - 支持双冒号(::)表示连续的零组（只能出现一次）
        - 支持IPv4映射的IPv6地址格式（如 ::ffff:192.168.1.1）

    :param ip: 待验证的IPv6地址字符串
    :return: 如果是有效的IPv6地址返回True，否则返回False
    """
    # 检查是否包含双冒号
    if "::" in ip:
        # 双冒号只能出现一次
        if ip.count("::") > 1:
            return False

        # 分割双冒号前后的部分
        parts = ip.split("::")
        left_parts = parts[0].split(":") if parts[0] else []
        right_parts = parts[1].split(":") if parts[1] else []

        # 去除空字符串
        left_parts = [p for p in left_parts if p]
        right_parts = [p for p in right_parts if p]

        # 总组数不能超过8
        if len(left_parts) + len(right_parts) >= 8:
            return False

        # 验证左右两部分
        all_parts = left_parts + right_parts
    else:
        # 没有双冒号的情况，必须恰好有8组
        all_parts = ip.split(":")
        if len(all_parts) != 8:
            return False

    # 验证每组是否为有效的十六进制数
    for part in all_parts:
        # 每组长度不能超过4个字符
        if len(part) > 4:
            return False

        # 检查是否为有效的十六进制数
        try:
            int(part, 16)
        except ValueError:
            # 特殊情况：检查是否为IPv4映射地址的最后一部分
            if "." in part:
                return is_ipv4(part)
            return False

    return True


class XdbSearcher(object):
    __f = None

    # 最小内存分配
    vector_index = None
    # 整个读取xdb，保存在内存中
    content_buff = None

    @staticmethod
    def load_vector_index_from_file(db_file) -> bytes | None:
        """
        从指定的db文件中加载向量索引数据到内存中

        :param db_file: IP地址库文件路径 (.xdb文件)
        :return: 向量索引数据的字节串，如果加载失败则返回 None

        :raises IOError: 当文件读取失败时记录异常日志
        """
        try:
            with open(db_file, "rb") as f:
                # 跳过文件头部信息区域，定位到向量索引开始位置
                f.seek(_HEADER_INFO_LENGTH)

                # 计算向量索引的总长度
                # 向量索引结构: 256行 × 256列 × 14字节(每个索引项大小)
                vi_len: int = _VECTOR_INDEX_ROWS * _VECTOR_INDEX_COLS * _SEGMENT_INDEX_SIZE

                # 读取向量索引数据
                vector_data: bytes = f.read(vi_len)
                return vector_data
        except IOError as e:
            logger.exception("加载ip2region.xdb文件向量索引数据到内存失败", e)
            return None

    @staticmethod
    def load_content_from_file(db_file) -> bytes | None:
        """
        从指定的db文件中加载全部数据到内存中

        :param db_file: IP地址库文件路径 (.xdb文件)
        :return: 文件的全部数据字节串，如果加载失败则返回 None

        :raises IOError: 当文件读取失败时记录异常日志
        """
        try:
            with open(db_file, "rb") as f:
                return f.read()
        except IOError as e:
            # 记录文件读取异常日志
            logger.exception("加载ip2region.xdb文件数据到内存失败", e)
            return None

    def __init__(self,
                 db_file: Optional[str] = None,
                 vector_index: Optional[bytes] = None,
                 content_buff: Optional[bytes] = None):
        """
        初始化 XdbSearcher 实例

        Note:
            - 如果提供 content_buff，则直接使用内存中的数据进行查询
            - 如果提供 vector_index，则使用内存中的向量索引进行快速查询
            - 如果都未提供，则通过文件句柄进行磁盘IO查询

        :param db_file: IP地址库文件路径 (.xdb文件)
        :param vector_index: 预加载的向量索引数据
        :param content_buff: 预加载的完整文件数据
        """
        self.init_database(db_file, vector_index, content_buff)

    def init_database(self,
                      db_file: Optional[str] = None,
                      vector_index: Optional[bytes] = None,
                      content_buff: Optional[bytes] = None):
        """
        初始化查询数据库

        :param db_file: IP地址库文件路径 (.xdb文件)
        :param vector_index: 预加载的向量索引数据
        :param content_buff: 预加载的完整文件数据
        """
        try:
            if content_buff is not None:
                self.__f = None
                self.vector_index = None
                self.content_buff = content_buff
            else:
                self.__f = open(db_file, "rb")
                self.vector_index = vector_index
        except IOError as e:
            logger.exception("初始化查询数据库失败", e)
            raise e

    def search(self, ip: str | int) -> str:
        """
        根据IP地址查询对应的地理位置信息

        Example:
           >>> searcher.search("192.168.1.1")
           "中国|0|广东省|深圳市|电信"
           >>> searcher.search(3232235777)
           "中国|0|广东省|深圳市|电信"

        :param ip: IPv4地址，可以是字符串格式("192.168.1.1")或整数格式
        :return: IP地址对应的地理位置信息，格式为"国家|区域|省份|城市|ISP"
                如果未找到匹配信息则返回空字符串
        """
        # 内网ipv6地址直接返回
        if ip == _IPV6_LOCAL_IP:
            return _IPV6_LOCAL_IP_REGION
        # ipv6地址不处理
        if is_ipv6(ip):
            return ""

        if isinstance(ip, str):
            return self.search_by_ip_str(ip)
        else:
            return self.search_by_ip_long(ip)

    def search_by_ip_str(self, ip: str) -> str:
        """
        根据IPv4地址字符串查询对应的地理位置信息

        Example:
            >>> searcher.search_by_ip_str("192.168.1.1")
            "中国|0|广东省|深圳市|电信"

        :param ip: IPv4地址字符串，例如 "192.168.1.1"
        :return: IP地址对应的地理位置信息，格式为"国家|区域|省份|城市|ISP"
                 如果未找到匹配信息则返回空字符串
        """
        # 内网ipv6地址直接返回
        if ip == _IPV6_LOCAL_IP:
            return _IPV6_LOCAL_IP_REGION
        # ipv6地址不处理
        if is_ipv6(ip):
            return ""

        # 检查字符串中的所有字符是否都是数字字符
        if not ip.isdigit():
            # 如果是"192.168.1.1"
            ip: int = ip_to_uint32(ip)
        return self.search_by_ip_long(int(ip))

    def search_by_ip_long(self, ip: int) -> str:
        """
        根据IPv4地址的整数形式查询对应的地理位置信息

        Note:
            - 使用二分查找算法在IP地址库中查找匹配的地理位置信息
            - 支持三种查询模式：向量索引、内存缓冲区、文件IO

        :param ip: IPv4地址的32位无符号整数表示
        :return: IP地址对应的地理位置信息，格式为"国家|区域|省份|城市|ISP"
                 如果未找到匹配信息则返回空字符串
        """
        # 根据向量索引定位段索引块
        start_ptr = end_ptr = 0

        # 计算IP地址的前两个字节
        ip_byte_0 = int((ip >> 24) & 0xFF)  # 第一个字节
        ip_byte_1 = int((ip >> 16) & 0xFF)  # 第二个字节
        # 计算向量索引位置
        vector_idx = ip_byte_0 * _VECTOR_INDEX_COLS * _VECTOR_INDEX_SIZE + ip_byte_1 * _VECTOR_INDEX_SIZE

        # 根据不同模式获取段索引块的起始和结束指针
        if self.vector_index is not None:
            # 使用预加载的向量索引
            start_ptr = get_uint32(self.vector_index, vector_idx)
            end_ptr = get_uint32(self.vector_index, vector_idx + 4)
        elif self.content_buff is not None:
            # 使用预加载的完整文件数据
            start_ptr = get_uint32(self.content_buff, _HEADER_INFO_LENGTH + vector_idx)
            end_ptr = get_uint32(self.content_buff, _HEADER_INFO_LENGTH + vector_idx + 4)
        else:
            # 使用文件句柄进行IO操作
            self.__f.seek(_HEADER_INFO_LENGTH + vector_idx)
            buffer_ptr = self.__f.read(8)
            start_ptr = get_uint32(buffer_ptr, 0)
            end_ptr = get_uint32(buffer_ptr, 4)

        # 在段索引块中使用二分查找算法查找地理位置信息
        data_length = data_pointer = int(-1)
        low = int(0)
        high = int((end_ptr - start_ptr) / _SEGMENT_INDEX_SIZE)

        while low <= high:
            mid = int((low + high) >> 1)
            position = int(start_ptr + mid * _SEGMENT_INDEX_SIZE)
            # 读取段索引
            buffer_sip = self.read_buffer(position, _SEGMENT_INDEX_SIZE)
            start_ip = get_uint32(buffer_sip, 0)

            if ip < start_ip:
                high = mid - 1
            else:
                end_ip = get_uint32(buffer_sip, 4)
                if ip > end_ip:
                    low = mid + 1
                else:
                    # 找到匹配的段索引，获取数据长度和指针
                    data_length = get_uint16(buffer_sip, 8)
                    data_pointer = get_uint32(buffer_sip, 10)
                    break

        # 空匹配拦截
        if data_pointer < 0:
            return ""

        # 读取并返回地理位置信息
        data_buffer = self.read_buffer(data_pointer, data_length)
        return data_buffer.decode("utf-8")

    def read_buffer(self, offset: int, length: int) -> bytes | None:
        """
        从内存缓冲区或文件中读取指定长度的字节数据

        Note:
           - 优先从内存缓冲区(content_buff)中读取数据
           - 如果内存缓冲区为空，则从文件句柄中读取数据
           - 该方法用于IP地址库查询过程中读取索引和数据块

        Example:
           >>> data = searcher.read_buffer(256, 100)
           >>> len(data)
           100

        :param offset: 读取数据的起始偏移位置
        :param length: 要读取的数据长度（字节数）
        :return: 读取到的字节数据，如果读取失败则返回 None
        """
        # 优先从内存缓冲区读取数据
        if self.content_buff is not None:
            return self.content_buff[offset:offset + length]

        # 如果内存缓存区为空，则从文件句柄中读取数据
        if self.__f is not None:
            self.__f.seek(offset)
            return self.__f.read(length)

        # 如果既没有内存缓冲区也没有文件句柄，返回 None
        return None

    def close(self):
        """
        关闭数据库文件句柄并清理资源

        Note:
            - 关闭打开的文件句柄
            - 清理向量索引和内容缓冲区引用
            - 该方法应在使用完查询器后调用以释放资源
        """
        if self.__f is not None:
            self.__f.close()
        self.vector_index = None
        self.content_buff = None


_searcher: XdbSearcher | None = None


def get_searcher() -> XdbSearcher:
    """
    获取全局唯一的 XdbSearcher 实例，用于查询 IP 地址对应的地理位置信息。

    Example:
        >>> searcher = get_searcher()
        >>> result = searcher.search("8.8.8.8")
        >>> print(result)
        "美国|0|加利福尼亚|洛杉矶|谷歌DNS"

    Note:
        - 该函数采用单例模式，确保全局只有一个 XdbSearcher 实例
        - 首次调用时会加载整个 IP 数据库文件到内存中以提高查询性能
        - 数据库文件路径为 TEMP_DIR 目录下的 ip2region.xdb 文件
        - 加载的数据会在内存中缓存，避免重复的文件 IO 操作
        - 适用于需要频繁进行 IP 查询的场景

    :return: 全局唯一的 XdbSearcher 实例，用于执行 IP 地理位置查询操作。
    :raises ValueError: 当无法从文件加载 IP 数据库内容时抛出异常。
    """
    global _searcher
    if _searcher is not None:
        return _searcher

    # 1. 缓存
    content_buff = XdbSearcher.load_content_from_file(db_file=f"{TEMP_DIR}/ip2region.xdb")
    if content_buff is None:
        raise IOError("Failed to load content from file")

    # 2. 创建查询对象
    _searcher = XdbSearcher(content_buff=content_buff)
    return _searcher


def close_searcher() -> None:
    """
    关闭全局 XdbSearcher 实例，释放资源。

    Note:
        - 该函数用于在程序退出时清理资源，确保 IP 地址查询器在程序结束后被正确关闭。
        - 关闭查询器会释放内存中的数据库内容，并关闭打开的文件句柄。
    """
    global _searcher
    if _searcher is not None:
        _searcher.close()
